#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
地理编码脚本 - 为map_data.json中的所有地址添加经纬度坐标
使用高德地图API批量获取坐标
"""

import json
import time
import requests
from typing import Dict, Tuple, Optional

# 高德地图API配置
AMAP_KEY = 'b2dec638f9508c079a4da1c7262ca67a'
GEOCODE_URL = 'https://restapi.amap.com/v3/geocode/geo'

class AmapGeocoder:
    """高德地图地理编码器"""
    
    def __init__(self, api_key: str):
        self.api_key = api_key
        self.success_count = 0
        self.fail_count = 0
        self.skip_count = 0
        
    def geocode(self, address: str) -> Optional[Tuple[float, float]]:
        """
        地理编码：将地址转换为经纬度
        
        Args:
            address: 地址字符串
            
        Returns:
            (经度, 纬度) 或 None
        """
        try:
            params = {
                'key': self.api_key,
                'address': address,
                'city': '北京市'
            }
            
            response = requests.get(GEOCODE_URL, params=params, timeout=10)
            response.raise_for_status()
            
            data = response.json()
            
            if data['status'] == '1' and data['count'] != '0':
                # 成功获取坐标
                location = data['geocodes'][0]['location']
                lng, lat = location.split(',')
                return (float(lng), float(lat))
            else:
                return None
                
        except Exception as e:
            print(f"  [ERROR] 请求异常: {e}")
            return None
    
    def process_address(self, name: str, address: str) -> Optional[Tuple[float, float]]:
        """
        处理单个地址
        
        Args:
            name: 名称（工作站或楼宇）
            address: 地址
            
        Returns:
            (经度, 纬度) 或 None
        """
        if not address or address == '':
            print(f"  [SKIP] 跳过（无地址）: {name}")
            self.skip_count += 1
            return None
        
        # 确保地址包含北京市
        full_address = address
        if '北京' not in address:
            if '石景山' in address:
                full_address = '北京市' + address
            else:
                full_address = '北京市石景山区' + address
        
        print(f"  正在编码: {name}")
        print(f"  地址: {full_address}")
        
        result = self.geocode(full_address)
        
        if result:
            lng, lat = result
            print(f"  [OK] 成功: {lng:.6f}, {lat:.6f}")
            self.success_count += 1
            return result
        else:
            print(f"  [FAIL] 失败")
            self.fail_count += 1
            return None


def process_map_data(data: Dict) -> Dict:
    """处理地图数据，为所有地址添加坐标"""
    
    geocoder = AmapGeocoder(AMAP_KEY)
    map_data = data['地图数据']
    
    total_items = 0
    processed_items = 0
    
    print("\n" + "="*60)
    print("开始处理地图数据...")
    print("="*60)
    
    # 遍历所有地图
    for map_key, map_info in map_data.items():
        print(f"\n处理地图: {map_info['地图名称']}")
        print("-" * 60)
        
        # 处理区级地图
        if map_info['地图类型'] == '区级':
            stations = map_info['数据'].get('工作站点位', [])
            
            for station in stations:
                total_items += 1
                
                # 检查是否已有坐标
                if station.get('经度') and station.get('纬度') and \
                   station['经度'] != '' and station['纬度'] != '':
                    print(f"  [EXIST] 已有坐标（工作站）: {station['工作站名称']}")
                else:
                    # 获取地址
                    address = station.get('完整地址') or station.get('工作站地址', '')
                    
                    # 地理编码
                    result = geocoder.process_address(station['工作站名称'], address)
                    
                    if result:
                        lng, lat = result
                        station['经度'] = lng
                        station['纬度'] = lat
                        processed_items += 1
                    
                    time.sleep(0.05)
                
                # 处理嵌套的覆盖楼宇列表
                buildings = station.get('覆盖楼宇列表', [])
                for building in buildings:
                    total_items += 1
                    
                    # 检查是否已有坐标
                    if building.get('经度') and building.get('纬度') and \
                       building['经度'] != '' and building['纬度'] != '':
                        print(f"    [EXIST] 已有坐标（楼宇）: {building['楼宇名称']}")
                    else:
                        # 获取地址
                        address = building.get('完整地址') or building.get('楼宇地址', '')
                        
                        # 地理编码
                        result = geocoder.process_address(building['楼宇名称'], address)
                        
                        if result:
                            lng, lat = result
                            building['经度'] = lng
                            building['纬度'] = lat
                            processed_items += 1
                        
                        time.sleep(0.05)
        
        # 处理街道/园区级地图
        else:
            # 处理工作站
            stations = map_info['数据'].get('工作站列表', [])
            
            for station in stations:
                total_items += 1
                
                # 检查是否已有坐标
                if station.get('经度') and station.get('纬度') and \
                   station['经度'] != '' and station['纬度'] != '':
                    print(f"  [EXIST] 已有坐标（工作站）: {station['工作站名称']}")
                else:
                    # 获取地址
                    address = station.get('工作站地址', '')
                    
                    # 地理编码
                    result = geocoder.process_address(station['工作站名称'], address)
                    
                    if result:
                        lng, lat = result
                        station['经度'] = lng
                        station['纬度'] = lat
                        processed_items += 1
                    
                    time.sleep(0.05)
                
                # 处理嵌套的覆盖楼宇列表
                nested_buildings = station.get('覆盖楼宇列表', [])
                for building in nested_buildings:
                    total_items += 1
                    
                    # 检查是否已有坐标
                    if building.get('经度') and building.get('纬度') and \
                       building['经度'] != '' and building['纬度'] != '':
                        print(f"    [EXIST] 已有坐标（嵌套楼宇）: {building['楼宇名称']}")
                    else:
                        # 获取地址
                        address = building.get('完整地址') or building.get('楼宇地址', '')
                        
                        # 地理编码
                        result = geocoder.process_address(building['楼宇名称'], address)
                        
                        if result:
                            lng, lat = result
                            building['经度'] = lng
                            building['纬度'] = lat
                            processed_items += 1
                        
                        time.sleep(0.05)
            
            # 处理楼宇
            buildings = map_info['数据'].get('楼宇列表', [])
            
            for building in buildings:
                total_items += 1
                
                # 检查是否已有坐标
                if building.get('经度') and building.get('纬度') and \
                   building['经度'] != '' and building['纬度'] != '':
                    print(f"  [EXIST] 已有坐标: {building['楼宇名称']}")
                    continue
                
                # 获取地址
                address = building.get('完整地址') or building.get('楼宇地址', '')
                
                # 地理编码
                result = geocoder.process_address(building['楼宇名称'], address)
                
                if result:
                    lng, lat = result
                    building['经度'] = lng
                    building['纬度'] = lat
                    processed_items += 1
                
                time.sleep(0.05)
    
    # 打印统计信息
    print("\n" + "="*60)
    print("[3/4] 处理完成！")
    print("="*60)
    print(f"总项目数: {total_items}")
    print(f"新增坐标: {processed_items}")
    print(f"[OK] 成功: {geocoder.success_count}")
    print(f"[FAIL] 失败: {geocoder.fail_count}")
    print(f"[SKIP] 跳过: {geocoder.skip_count}")
    print(f"[EXIST] 已有坐标: {total_items - processed_items - geocoder.fail_count - geocoder.skip_count}")
    print("="*60)
    
    return data


def main():
    """主函数"""
    
    print("\n高德地图地理编码工具")
    print("=" * 60)
    
    # 读取数据
    print("\n[1/4] 读取 map_data.json...")
    try:
        with open('map_data.json', 'r', encoding='utf-8') as f:
            data = json.load(f)
        print(f"成功读取，共 {data['地图数量']} 张地图")
    except Exception as e:
        print(f"读取失败: {e}")
        return
    
    # 备份原文件
    print("\n[2/4] 备份原文件...")
    try:
        with open('map_data.json.backup', 'w', encoding='utf-8') as f:
            json.dump(data, f, ensure_ascii=False, indent=2)
        print("已备份到 map_data.json.backup")
    except Exception as e:
        print(f"备份失败: {e}")
    
    # 处理数据
    processed_data = process_map_data(data)
    
    # 保存结果
    print("\n[4/4] 保存处理后的数据...")
    try:
        with open('map_data.json', 'w', encoding='utf-8') as f:
            json.dump(processed_data, f, ensure_ascii=False, indent=2)
        print("已保存到 map_data.json")
    except Exception as e:
        print(f"保存失败: {e}")
        return
    
    print("\n全部完成！现在可以刷新网页查看地图了！")


if __name__ == '__main__':
    main()

